package com.cmbchina.encrypt;

import org.bouncycastle.jce.provider.BouncyCastleProvider;

import javax.crypto.Cipher;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.security.Provider;
import java.security.Security;

/**
 * SM4
 *
 * @author liuxb
 */
public class SM4 {
    public static final Provider PROVIDER = new BouncyCastleProvider();
    private static final String SM4 = "SM4";
    private static final int BLOCK_SIZE = 16;
    private final PKCS5Padding padding = new PKCS5Padding(16);

    public SM4() {
    }

    public static SM4 getInstance() {
        Security.addProvider(PROVIDER);
        return new SM4();
    }

    private byte[] doECBCipher(byte[] key, byte[] input, int operateMode) throws SMCryptException {
        if (key != null && key.length == BLOCK_SIZE) {
            if (input == null) {
                throw new SMCryptKYException(ErrorCode.E10400);
            } else {
                if (operateMode == 1) {
                    Checker.check(input != null && input.length > 0, ErrorCode.E10408);
                    input = this.padding.pad(input);
                } else {
                    Checker.check(input != null && input.length > 0 && input.length % BLOCK_SIZE == 0, ErrorCode.E10409);
                }
                SecretKeySpec spec = new SecretKeySpec(key, "SM4");
                try {
                    Cipher cipher = Cipher.getInstance("SM4/ECB/NoPadding", PROVIDER);
                    cipher.init(operateMode, spec);
                    return operateMode == 1 ? cipher.doFinal(input) : this.padding.unpad(cipher.doFinal(input));
                } catch (Exception var6) {
                    if (operateMode == 1) {
                        throw new SMCryptKYException(ErrorCode.E10204, var6);
                    } else {
                        throw new SMCryptKYException(ErrorCode.E10205, var6);
                    }
                }
            }
        } else {
            throw new SMCryptKYException(ErrorCode.E10410);
        }
    }

    private byte[] doCipher(byte[] key, byte[] iv, byte[] input, String mode, int operateMode) throws SMCryptException {
        if (key != null && key.length == BLOCK_SIZE) {
            if (input == null) {
                throw new SMCryptKYException(ErrorCode.E10400);
            } else if (iv == null) {
                throw new SMCryptKYException(ErrorCode.E10411);
            } else {
                if (operateMode == 1) {
                    if (input == null || input.length <= 0) {
                        throw new SMCryptKYException(ErrorCode.E10408);
                    }
                    input = this.padding.pad(input);
                } else if (input == null || input.length <= 0 || input.length % BLOCK_SIZE != 0) {
                    throw new SMCryptKYException(ErrorCode.E10409);
                }
                if (iv != null && iv.length == BLOCK_SIZE) {
                    SecretKeySpec spec = new SecretKeySpec(key, "SM4");
                    IvParameterSpec ivParameterSpec = new IvParameterSpec(iv);
                    try {
                        Cipher cipher = Cipher.getInstance("SM4/" + mode + "/NoPadding", PROVIDER);
                        cipher.init(operateMode, spec, ivParameterSpec);
                        if (operateMode == 1) {
                            return cipher.doFinal(input);
                        } else {
                            return this.padding.unpad(cipher.doFinal(input));
                        }
                    } catch (Exception var9) {
                        if (operateMode == 1) {
                            throw new SMCryptKYException(ErrorCode.E10204, var9);
                        } else {
                            throw new SMCryptKYException(ErrorCode.E10205, var9);
                        }
                    }
                } else {
                    throw new SMCryptKYException(ErrorCode.E10411);
                }
            }
        } else {
            throw new SMCryptKYException(ErrorCode.E10410);
        }
    }

    public byte[] CMBSM4EncryptWithECB(byte[] key, byte[] input) throws SMCryptException {
        return this.doECBCipher(key, input, 1);
    }

    public byte[] CMBSM4DecryptWithECB(byte[] key, byte[] input) throws SMCryptException {
        return this.doECBCipher(key, input, 2);
    }

    public byte[] CMBSM4EncryptWithCBC(byte[] key, byte[] iv, byte[] input) throws SMCryptException {
        return this.doCipher(key, iv, input, "CBC", 1);
    }

    public byte[] CMBSM4DecryptWithCBC(byte[] key, byte[] iv, byte[] input) throws SMCryptException {
        return this.doCipher(key, iv, input, "CBC", 2);
    }

    public byte[] CMBSM4EncryptWithCFB(byte[] key, byte[] iv, byte[] input) throws SMCryptException {
        return this.doCipher(key, iv, input, "CFB", 1);
    }

    public byte[] CMBSM4DecryptWithCFB(byte[] key, byte[] iv, byte[] input) throws SMCryptException {
        return this.doCipher(key, iv, input, "CFB", 2);
    }

    public byte[] CMBSM4EncryptWithOFB(byte[] key, byte[] iv, byte[] input) throws SMCryptException {
        return this.doCipher(key, iv, input, "OFB", 1);
    }

    public byte[] CMBSM4DecryptWithOFB(byte[] key, byte[] iv, byte[] input) throws SMCryptException {
        return this.doCipher(key, iv, input, "OFB", 2);
    }

    public byte[] CMBSM4EncryptWithCTR(byte[] key, byte[] input) throws SMCryptException {
        return this.doCipher(key, new byte[16], input, "CTR", 1);
    }

    public byte[] CMBSM4DecryptWithCTR(byte[] key, byte[] input) throws SMCryptException {
        return this.doCipher(key, new byte[16], input, "CTR", 2);
    }

    public byte[] CMBSM4EncryptWithCTR(byte[] key, byte[] iv, byte[] input) throws SMCryptException {
        return this.doCipher(key, iv, input, "CTR", 1);
    }

    public byte[] CMBSM4DecryptWithCTR(byte[] key, byte[] iv, byte[] input) throws SMCryptException {
        return this.doCipher(key, iv, input, "CTR", 2);
    }
}